import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.io.*;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.*;

@ServerEndpoint("/")
public class WebSocketServer {
    
    private static final int PORT = 9000;
    private static final String AUDIO_DIR = "audio/";
    private static final String FILE_TO_PLAY = "vacation.wav";
    private static final String CONVERTED_FILE = "audio/converted.wav";
    
    private Session session;
    private String username = "unnamed client";
    private String uuid = UUID.randomUUID().toString();
    private boolean responseType = false;
    private int recvPackets = 0;
    private int sendPackets = 0;
    private int packetLen = 800;
    
    private static final Map<Session, WebSocketServer> clients = new ConcurrentHashMap<>();
    private ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
    private ScheduledFuture<?> heartbeatFuture;
    // thread pool
    private static final ExecutorService audioExecutor = Executors.newSingleThreadExecutor();
    
    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
        clients.put(session, this);
        System.out.println("客户端连接: " + session.getId());
        
        // 启动心跳检测
        startHeartbeat();
    }
    
    @OnMessage
    public void onMessage(String message, Session session) {
        System.out.println("收到文本消息: " + message);
        
        try {
            // 简单JSON解析
            if (message.contains("\"method\":\"start\"")) {
                handleStart(message);
            } else if (message.contains("\"method\":\"stop\"")) {
                handleStop();
            } else if (message.contains("\"method\":\"ping\"")) {
                handlePing();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    
    @OnMessage
    public void onMessage(ByteBuffer message, Session session) {
        recvPackets++;
        System.out.println("收到二进制数据, 长度: " + message.remaining());
    }
    
    @OnClose
    public void onClose(Session session) {
        audioExecutor.shutdownNow();
        clients.remove(session);
        System.out.println("客户端断开: " + session.getId());
        stopHeartbeat();
    }
    
    @OnError
    public void onError(Session session, Throwable error) {
        error.printStackTrace();
        clients.remove(session);
        System.out.println("客户端错误: " + session.getId());
        stopHeartbeat();
    }
    
    private void handleStart(String message) {
        // 解析参数
        String audio = message.contains("\"audio\":\"sendonly\"") ? "sendonly" : 
                      message.contains("\"audio\":\"recvonly\"") ? "recvonly" : 
                      message.contains("\"audio\":\"inactive\"") ? "inactive" : "sendrecv";
        
        responseType = !audio.equals("sendonly");
        
        // 发送响应
        String response = "{\"jsonrpc\":\"2.0\",\"id\":0,\"result\":{\"code\":200,\"message\":\"OK\",\"audio\":\"" + 
                         audio + "\",\"heartbeat\":100,\"mux_type\":\"mux\"}}";
        
        try {
            session.getBasicRemote().sendText(response);
            
            // 如果需要发送音频
            if (responseType) {
                playAudioFile();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
    
    private void handleStop() {
        System.out.println("处理停止请求");
        try {
            session.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
    
    private void handlePing() {
        String pong = "{\"jsonrpc\":\"2.0\",\"method\":\"pong\",\"result\":\"server is alive.\"}";
        try {
            session.getBasicRemote().sendText(pong);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
    
    private void playAudioFile() {
        audioExecutor.submit(() -> {
            try {
                // 1. 读取原始音频文件
                byte[] originalData = Files.readAllBytes(Paths.get(AUDIO_DIR + FILE_TO_PLAY));
                System.out.println("原始音频大小: " + originalData.length + " 字节");

                // 2. 转换为8000Hz, 16bit, 单声道
                byte[] convertedData = WavConverter.convertWav(originalData);
                System.out.println("转换后音频大小: " + convertedData.length + " 字节");

                // 3. 保存转换后的文件
                Files.write(Paths.get(AUDIO_DIR + "converted.wav"), convertedData);

                // 4. 分片发送音频数据
                int chunkSize = packetLen * 2;
                for (int i = 0; i < convertedData.length; i += chunkSize) {
                    if (!session.isOpen()) break;
                    
                    int end = Math.min(i + chunkSize, convertedData.length);
                    byte[] chunk = Arrays.copyOfRange(convertedData, i, end);
                    
                    // 异步发送音频片段
                    session.getAsyncRemote().sendBinary(ByteBuffer.wrap(chunk));
                    sendPackets++;
                    System.out.printf("发送音频片段 %d/%d (%.1f%%) 长度: %d%n",
                        i/chunkSize, 
                        convertedData.length/chunkSize,
                        (i*100.0/convertedData.length),
                        chunk.length);
                    
                    Thread.sleep(20); // 保持发送间隔
                }

                // 5. 发送停止消息
                // 可选，发送此消息后将停止放音
                if (session.isOpen()) {
                    String stopMsg = "{\"jsonrpc\":\"2.0\",\"method\":\"stop_play_audio\"}";
                    // session.getAsyncRemote().sendText(stopMsg);
                }

            } catch (Exception e) {
                System.err.println("音频发送出错: ");
                e.printStackTrace();
            }
        });
    }
    
    private void startHeartbeat() {
        heartbeatFuture = scheduler.scheduleAtFixedRate(() -> {
            try {
                if (session.isOpen()) {
                    session.getBasicRemote().sendText("{\"jsonrpc\":\"2.0\",\"method\":\"ping\"}");
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }, 5, 5, TimeUnit.SECONDS);
    }
    
    private void stopHeartbeat() {
        if (heartbeatFuture != null) {
            heartbeatFuture.cancel(true);
        }
    }
    
    public static void main(String[] args) {
        // 检查音频目录是否存在
        File audioDir = new File(AUDIO_DIR);
        if (!audioDir.exists()) {
            audioDir.mkdir();
            System.out.println("已创建音频目录: " + AUDIO_DIR);
        }
        
        // 使用内置Jetty服务器
        org.eclipse.jetty.server.Server server = new org.eclipse.jetty.server.Server(PORT);
        
        // 创建WebSocket处理器
        org.eclipse.jetty.servlet.ServletContextHandler context = 
            new org.eclipse.jetty.servlet.ServletContextHandler(
                org.eclipse.jetty.servlet.ServletContextHandler.SESSIONS);
        context.setContextPath("/");
        server.setHandler(context);
        
        // 添加WebSocket端点
        org.eclipse.jetty.websocket.jsr356.server.deploy.WebSocketServerContainerInitializer.configure(
            context, (servletContext, wsContainer) -> {
                wsContainer.addEndpoint(WebSocketServer.class);
            });
        
        try {
            server.start();
            System.out.println("WebSocket服务器启动，端口: " + PORT);
            server.join();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
